#include "colorVocabulary.h"
#include "grabcut.h"
#include "configurationParameters.h"
#include "utilitiesFunction.h"
#include <iostream>
#include <string>

colorVocabulary::colorVocabulary(int dictionarySize, int sampleSize) {
	this->dictionarySize = dictionarySize;
	this->sampleSize = sampleSize;
}

void colorVocabulary::getSampleOnImage(char *filename, int numSample) {
	cv::Mat img = cv::imread(std::string(filename));
	// Convert RGB to Lab color
	cv::Mat imgLAB;
	cvtColor(img, imgLAB, CV_BGR2Lab);
	Grabcut grabcut;
	grabcut.segment(img, configurationParameters::gcRatioCenter());
	cv::Mat grabcutMask = grabcut.loadMask();
	std::cout << "Finish GC" << std::endl;

	// Dense Keypoint
	float initFeatureScale = 1.f;
	int featureScaleLevels = 1;
	float featureScaleMul = 0.1f;
	int initXyStep = 6;
	int initImgBound = 0;
	bool varyXyStepWithScale = true;
	bool varyImgBoundWithScale = false;

	DenseFeatureDetector detector(initFeatureScale, featureScaleLevels, featureScaleMul, initXyStep, initImgBound, varyXyStepWithScale, varyImgBoundWithScale);
	vector<KeyPoint> denseQueryKeypoints;
	detector.detect(img, denseQueryKeypoints, grabcutMask);

	// Create random sample index keypoints on image
	int max = denseQueryKeypoints.size() - 1;
	std::cout << "denseQueryKeypoints.size(): " << denseQueryKeypoints.size() << std::endl;

	vector<int> index_random;

	if (numSample > denseQueryKeypoints.size()){
		numSample = denseQueryKeypoints.size();
		for (int i = 0; i <denseQueryKeypoints.size(); i++){
			index_random.push_back(i);
		}
	}
	else{
		int random;
		int uniqueflag;
		int i, j;

		for (i = 0; i < numSample; i++) {
			do {
				/* Assume things are unique... we'll reset this flag if not. */
				uniqueflag = 1;
				random = rand() % max + 1;
				/* This loop checks for uniqueness */
				for (j = 0; j < i && uniqueflag == 1; j++) {
					if (index_random.at(j) == random) {
						uniqueflag = 0;
					}
				}
			} while (uniqueflag != 1);
			index_random.push_back(random);
		}
	}

	std::cout << "Create random sample index keypoints on image" << std::endl;
	std::cout << "index_random.size(): " << index_random.size() << std::endl;
	std::cout << "number_sample_on_image: " << numSample << std::endl;

	// Get sample keypoints
	vector<KeyPoint> sample_keypoint_on_image;

	for (int i = 0; i < numSample; i++) {
		int index_keypoint = index_random.at(i);
		try{
			KeyPoint kp = denseQueryKeypoints.at(index_keypoint);
			sample_keypoint_on_image.push_back(kp);
			this->sampleKeypoint.push_back(kp);
		}
		catch (const std::exception& ex){
			std::cout << "Error at: " << index_keypoint << std::endl;
			int stop;
			std::cin >> stop;
		}

	}

	std::cout << "Sample keypoint on image: " << sample_keypoint_on_image.size() << std::endl;
	std::cout << "Get sample keypoints" << std::endl;

	// Get sample color descriptor on image
	int numRow = sample_keypoint_on_image.size();
	std::cout << "Num Row = " << numRow << std::endl;
	cv::Mat listDescriptor;
	for (int i = 0; i < numRow; ++i) {
		cv::KeyPoint tempKeyPoint = sample_keypoint_on_image.at(i);
		int x = tempKeyPoint.pt.x;
		int y = tempKeyPoint.pt.y;
		cv::Vec3f descriptor = imgLAB.ptr<Vec3f>(y)[x];

		double L = descriptor.val[0];
		double a = descriptor.val[1];
		double b = descriptor.val[2];
		cv::Mat row = cv::Mat(1, 3, CV_32FC1, { L, a, b });
		listDescriptor.push_back(row);
	}

	std::cout << "Get sample sift descriptor on image" << std::endl;

	sampleColorDescriptor.push_back(listDescriptor);

	std::cout << sampleColorDescriptor.size() << std::endl;
	std::cout << std::endl;
	std::cout << "====================" << std::endl;
}

void colorVocabulary::getSampleOnTrainingSet(std::string directory, std::string fileNameFormat, int fromFile, int toFile) {
	int number_files = (toFile - fromFile) + 1;
	int number_sample_on_image = this->sampleSize / number_files;
	//int number_sample_on_image = 100;
	int re_assign_number_sample = number_sample_on_image * number_files;
	this->sampleSize = re_assign_number_sample;

	// Call getSampleOnImage
	for (int i = fromFile; i <= toFile; i++){
		string order_file = "";
		if (i < 10){
			order_file = "000" + std::to_string(i);

		}
		if (10 <= i & i < 100){
			order_file = "00" + std::to_string(i);

		}
		if (100 <= i & i < 1000){
			order_file = "0" + std::to_string(i);
		}
		if (1000 <= i & i < 10000){
			order_file = std::to_string(i);
		}

		string string_path_file_image = directory + fileNameFormat + "_" + order_file + ".jpg";

		std::cout << string_path_file_image << std::endl;
		std::cout << "====================" << std::endl;
		std::cout << std::endl;
		char *cstr_image = new char[string_path_file_image.length() + 1];
		strcpy(cstr_image, string_path_file_image.c_str());
		// do stuff
		getSampleOnImage(cstr_image, number_sample_on_image);
		delete[] cstr_image;
	}
	this->sampleSize = this->sampleColorDescriptor.rows;
	std::cout << "Real sample size: " << this->sampleColorDescriptor.size() << std::endl;
	std::cout << "Rows size: " << this->sampleColorDescriptor.rows << std::endl;
	std::cout << "Cols size: " << this->sampleColorDescriptor.cols << std::endl;
	std::cout << "Finish to get sample!" << std::endl;
}

void colorVocabulary::generateDictionary(){
	utilitiesFunction::writeXML(sampleColorDescriptor, "sampleColorDescriptor.xml");
	
	cv::TermCriteria tc(CV_TERMCRIT_ITER, 100, 0.001);
	int retries = 1;
	int flags = KMEANS_PP_CENTERS;
	std::cout << "Dictionary size = " << dictionarySize << std::endl;
	cv::BOWKMeansTrainer bowTrainer(this->dictionarySize, tc, retries, flags);

	std::cout << "Real sample size: " << this->sampleColorDescriptor.size() << std::endl;
	std::cout << "Rows size: " << this->sampleColorDescriptor.rows << std::endl;
	std::cout << "Cols size: " << this->sampleColorDescriptor.cols << std::endl;

	std::cout << "Begining cluster!" << std::endl;
	this->dictionary = bowTrainer.cluster(this->sampleColorDescriptor);
	std::cout << "Finishing cluster!" << std::endl;
}

int colorVocabulary::getSizeDictionary(){
	return this->dictionarySize;
}

Mat colorVocabulary::getDictionary(){
	return this->dictionary;
}